from torch import nn

class MobileNetV1(nn.Module):
    # 用于控制总参数量
    alpha = 0.75
    beta = 1
    def __init__(self, classes):
        super().__init__()
        self.classes = classes
        self.feature, self.avg, self.fc = self.__create_nn()

    def __create_nn(self):
        layer = nn.Sequential()
        layer.add_module("conv1", self._Conv2DBlock(3, int(32 * self.alpha), 3, 2, 1))
        layer.add_module("conv2", self._Conv2DBlock(int(32 * self.alpha), int(32 * self.alpha), 3, 1, 1, int(32 * self.alpha)))
        layer.add_module("conv3", self._Conv2DBlock(int(32 * self.alpha), int(64 * self.alpha), 1, 1))
        layer.add_module("conv4", self._Conv2DBlock(int(64 * self.alpha), int(64 * self.alpha), 3, 2, 1, int(64 * self.alpha)))
        layer.add_module("conv5", self._Conv2DBlock(int(64 * self.alpha), int(128 * self.alpha), 1, 1))
        layer.add_module("conv6", self._Conv2DBlock(int(128 * self.alpha), int(128 * self.alpha), 3, 1, 1, int(128 * self.alpha)))
        layer.add_module("conv7", self._Conv2DBlock(int(128 * self.alpha), int(128 * self.alpha), 3, 1, 1))
        layer.add_module("conv8", self._Conv2DBlock(int(128 * self.alpha), int(128 * self.alpha), 3, 2, 1, int(128 * self.alpha)))
        layer.add_module("conv9", self._Conv2DBlock(int(128 * self.alpha), int(256 * self.alpha), 3, 1, 1))
        layer.add_module("conv10", self._Conv2DBlock(int(256 * self.alpha), int(256 * self.alpha), 3, 1, 1, int(256 * self.alpha)))
        layer.add_module("conv11", self._Conv2DBlock(int(256 * self.alpha), int(256 * self.alpha), 1, 1))
        layer.add_module("conv12", self._Conv2DBlock(int(256 * self.alpha), int(256 * self.alpha), 3, 2, 1, int(256 * self.alpha)))
        layer.add_module("conv13", self._Conv2DBlock(int(256 * self.alpha), int(512 * self.alpha), 1, 1))
        layer.add_module("conv_block", self._Conv2DBlock5x())
        layer.add_module("conv14", self._Conv2DBlock(int(512 * self.alpha), int(512 * self.alpha), 3, 2, 1, int(512 * self.alpha)))
        layer.add_module("conv15", self._Conv2DBlock(int(512 * self.alpha), int(1024 * self.alpha), 1, 1))
        layer.add_module("conv16", self._Conv2DBlock(int(1024 * self.alpha), int(1024 * self.alpha), 3, 1, 1, int(1024 * self.alpha)))
        layer.add_module("conv17", self._Conv2DBlock(int(1024 * self.alpha), int(1024 * self.alpha), 1, 1))
        # layer.add_module("avg", nn.AvgPool2d(7))
        # layer.add_module("fc", nn.Linear(int(1024 * self.alpha), self.classes))
        # return layer
        avg_pooling = nn.AvgPool2d(7)
        fc = nn.Linear(int(1024 * self.alpha), self.classes)
        return layer, avg_pooling, fc

    def _Conv2DBlock(self, in_c, out_c, k, s, p=0, groups=1):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, k, s, p, groups=groups, bias=False),
            nn.BatchNorm2d(out_c),
            nn.ReLU6(True)
        )

    def _Conv2DBlock5x(self):
        layer = []
        for i in range(5):
            layer.append(self._Conv2DBlock(int(512 * self.alpha), int(512 * self.alpha), 3, 1, 1, groups=int(512 * self.alpha)))
            layer.append(self._Conv2DBlock(int(512 * self.alpha), int(512 * self.alpha), 1, 1))
        return nn.Sequential(*layer)

    def forward(self, x):
        out = self.feature(x)
        out = self.avg(out).view(out.size(0), -1)
        out = self.fc(out)
        return out

def mobilenetv1():
    return MobileNetV1(5)

if __name__ == "__main__":
    from utils.utils import get_model_flops_args
    model = MobileNetV1(5)
    get_model_flops_args(model, (3, 224, 224))

